#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
键盘处理模块 - 用于处理远程控制过程中的键盘输入
支持录制控制、显示统计信息和机械臂控制等功能
"""

import cv2
import logging
import sys
import os

# 添加项目根目录到路径
current_dir = os.path.dirname(os.path.abspath(__file__))
parent_dir = os.path.dirname(os.path.dirname(current_dir))
sys.path.append(parent_dir)

# 导入日志模块
from teleop.utils.logger import get_logger

class KeyboardHandler:
    """
    键盘处理类，用于处理远程控制过程中的键盘输入
    """
    def __init__(self, arm_ctrl=None, recorder=None, record_enabled=False):
        """
        初始化键盘处理器
        
        参数:
            arm_ctrl: 机械臂控制器对象
            recorder: 录制器对象
            record_enabled: 是否启用录制功能
        """
        self.logger = get_logger('KeyboardHandler')
        self.arm_ctrl = arm_ctrl
        self.recorder = recorder
        self.record_enabled = record_enabled
        self.recording = False
        self.running = True
        
    def process_keyboard_input(self, image):
        """
        处理键盘输入并显示图像
        
        参数:
            image: 要显示的图像
            
        返回:
            bool: 是否继续运行程序
        """
        # 显示图像前添加状态信息
        self.add_status_info(image)
        
        # 显示图像
        cv2.imshow("record image", image)
        
        # 处理键盘输入 - 增加等待时间以提高按键捕获率
        key = cv2.waitKey(10) & 0xFF
        
        if key == ord('q'):
            self.running = False
        elif key == ord('s') and self.record_enabled and self.recorder:
            self._handle_record_toggle()
        elif key == ord('i') and self.record_enabled and self.recorder:
            # 显示录制器统计信息
            self.recorder.print_stats()
        elif key == ord('h') and self.arm_ctrl:
            # 按h键回到初始位置
            self.logger.info("键盘命令: 回到初始位置")
            self.arm_ctrl.go_home()
            
        return self.running
    
    def _handle_record_toggle(self):
        """处理录制状态切换"""
        if self.recording:
            # 如果正在录制，结束录制
            self.logger.info("结束记录......")
            # 使用新的 save_episode 接口，可以指定等待超时时间
            if self.recorder.save_episode(wait_timeout=2.0):  # 等待最多2秒
                self.logger.info("保存请求已发送，后台处理中")
            else:
                self.logger.info("保存请求发送失败")
            self.recording = False
        else:
            # 如果未录制，尝试开始新录制
            self.logger.info("尝试开始录制......")
            # 使用新的 create_episode 接口，可以指定等待超时时间
            if self.recorder.create_episode(timeout=3.0):  # 等待最多3秒
                self.recording = True
                self.logger.info(f"开始录制成功，episode_id: {self.recorder.episode_id}")
            else:
                self.logger.info("创建录制失败，系统可能正忙")
                self.recording = False
    
    def is_recording(self):
        """获取当前录制状态"""
        return self.recording
        
    def add_status_info(self, image):
        """
        在图像上添加状态信息
        
        参数:
            image: 要添加状态信息的图像
        """
        # 添加录制状态信息
        status_text = f"recording: {'on' if self.recording else 'off'}"
        cv2.putText(image, status_text, (10, 30), 
                    cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 0, 255), 2)
        
        # 显示保存状态
        if self.recorder and not self.recorder.is_available:
            cv2.putText(image, "saving...", (10, 60), 
                        cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 0, 255), 2)
